package com.ezandroid.alphago.engine;

import android.content.Context;

import org.tensorflow.contrib.android.TensorFlowInferenceInterface;

/**
 * 策略网络
 *
 * @author like
 * @date 2017-07-17
 */
public class PolicyNN {

    private static final int FEATURE_LENGTH = 48;
    private static final int BOARD_SIZE = 19;
    private static final String INPUT_NAME = "convolution2d_input_1";
    private static final String MODEL_FILE = "rocalphaV108P5175.pb";
    private static final String OUTPUT_NAME = "Softmax";
    private static String[] mOutputNames;
    private static TensorFlowInferenceInterface mTensorFlow;

    public PolicyNN(Context context) {
        mTensorFlow = new TensorFlowInferenceInterface(context.getAssets(), MODEL_FILE);
        mOutputNames = new String[]{OUTPUT_NAME};
    }

    public float[] getOutput(byte[][] input) {
        float[] in = new float[FEATURE_LENGTH * BOARD_SIZE * BOARD_SIZE];
        float[] output = new float[BOARD_SIZE * BOARD_SIZE];
        for (int i = 0; i < BOARD_SIZE; i++) {
            for (int j = 0; j < BOARD_SIZE; j++) {
                for (int k = 0; k < FEATURE_LENGTH; k++) {
                    in[j * BOARD_SIZE + i + k * BOARD_SIZE * BOARD_SIZE] =
                            input[j * BOARD_SIZE + i][k];
                }
            }
        }
        mTensorFlow.feed(INPUT_NAME, in, 1L, FEATURE_LENGTH, BOARD_SIZE, BOARD_SIZE);
        mTensorFlow.run(mOutputNames);
        mTensorFlow.fetch(OUTPUT_NAME, output);
        return output;
    }
}
